{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "zMxsnhXN_7rt"
      },
      "source": [
        "# LightGlue-ONNX Evaluation\n",
        "\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/fabio-sim/LightGlue-ONNX/blob/main/evaluation/lightglue-onnx.ipynb) [![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://raw.githubusercontent.com/fabio-sim/LightGlue-ONNX/main/evaluation/lightglue-onnx.ipynb)\n",
        "\n",
        "The inference time of [LightGlue-ONNX](https://github.com/fabio-sim/LightGlue-ONNX) is measured and compared to the [original PyTorch implementation](https://github.com/cvg/LightGlue).\n",
        "\n",
        "## Install Dependencies\n",
        "\n",
        "Ensure that you are connected to a GPU Runtime. Otherwise, go to `Runtime` -> `Change runtime type` and select the GPU hardware accelerator."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VKnZ2qk0_1YV",
        "outputId": "45af6d39-65ee-49ce-ba57-424168e5d51d"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi\n",
        "!nvcc --version\n",
        "!lsb_release -a"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "braER_YwBUHl",
        "outputId": "c7edda78-6d03-488e-9727-1c5d71d75770"
      },
      "outputs": [],
      "source": [
        "!pip install -q torch==2.0.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118\n",
        "!pip install -q kornia==0.6.12 onnx==1.14.0 onnxruntime-gpu==1.15.1 numpy opencv-python matplotlib einops\n",
        "# !pip install -q flash-attn==1.0.8 --no-build-isolation  # Time-consuming (~30 minutes)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Install TensorRT\n",
        "!wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/8.6.1/local_repos/nv-tensorrt-local-repo-ubuntu2204-8.6.1-cuda-11.8_1.0-1_amd64.deb\n",
        "!sudo dpkg -i nv-tensorrt-local-repo-ubuntu2204-8.6.1-cuda-11.8_1.0-1_amd64.deb\n",
        "!sudo cp /var/nv-tensorrt-local-repo-ubuntu2204-8.6.1-cuda-11.8/*-keyring.gpg /usr/share/keyrings/\n",
        "!sudo cp /var/nv-tensorrt-local-repo-ubuntu2204-8.6.1-cuda-11.8/nv-tensorrt-local-0628887B-keyring.gpg /usr/share/keyrings/\n",
        "!sudo apt-get update\n",
        "!sudo apt-get install tensorrt python3-libnvinfer-dev\n",
        "!dpkg-query -W tensorrt"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "tFAJLnp0Cewt"
      },
      "source": [
        "## Code & Data\n",
        "\n",
        "Clone the repository for the evaluation script."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "g0dSlABBBzJs",
        "outputId": "02050c34-abab-4af2-b073-b9d604486d20"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/fabio-sim/LightGlue-ONNX"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ESWw9UmUDa5m"
      },
      "source": [
        "The [MegaDepth test set](https://drive.google.com/file/d/12yKniNWebDHRTCwhBNJmxYMPgqYX3Nhv/view?usp=drive_link) (403 image pairs) provided by [LoFTR](https://github.com/zju3dv/LoFTR/blob/master/docs/TRAINING.md) will be used for measuring the inference times. Download the data and extract."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Qnn8kchyDWRG",
        "outputId": "f1564b12-5515-4055-e110-28724c83f440"
      },
      "outputs": [],
      "source": [
        "!cd LightGlue-ONNX && gdown 12yKniNWebDHRTCwhBNJmxYMPgqYX3Nhv\n",
        "!cd LightGlue-ONNX && tar -xvf megadepth_test_1500.tar"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "oQQO9xeXEn4F"
      },
      "source": [
        "## ONNX Export\n",
        "\n",
        "If you haven't exported any ONNX models, please run the export script first. Otherwise, please upload the models to the local `LightGlue-ONNX/weights` folder with the filename formatted as `{extractor_type}_{max_num_keypoints}.onnx`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IeAIIwcCE5UV",
        "outputId": "b71b9a1f-525e-4086-97f6-3e770fa8f995"
      },
      "outputs": [],
      "source": [
        "!cd LightGlue-ONNX && python export.py \\\n",
        "  --extractor_type superpoint \\\n",
        "  --extractor_path weights/superpoint_512.onnx \\\n",
        "  --dynamic --max_num_keypoints 512"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# For TensorRT\n",
        "!cd LightGlue-ONNX && python tools/symbolic_shape_infer.py \\\n",
        "  --input weights/superpoint_lightglue.onnx \\\n",
        "  --output weights/superpoint_lightglue.onnx \\\n",
        "  --auto_merge"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "dMUNFTppEX8d"
      },
      "source": [
        "## Evaluation\n",
        "\n",
        "Measure the average inference time of the original PyTorch implementation with the default configuration for LightGlue."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aB2nOT8hDW_o",
        "outputId": "6763dde7-3d48-475b-e2e0-e3b496ee653b"
      },
      "outputs": [],
      "source": [
        "!cd LightGlue-ONNX && python eval.py torch \\\n",
        "  --img_size 1024 \\\n",
        "  --extractor_type superpoint \\\n",
        "  --max_num_keypoints 512 \\\n",
        "  --device cuda"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V6MCEYDEGIZm",
        "outputId": "808964a0-9d45-4b94-be16-df96d5179caf"
      },
      "outputs": [],
      "source": [
        "!cd LightGlue-ONNX && python eval.py ort \\\n",
        "  --img_size 1024 \\\n",
        "  --extractor_type superpoint \\\n",
        "  --max_num_keypoints 512 \\\n",
        "  --device cuda"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "CFjlC9LVHsnf"
      },
      "source": [
        "# References\n",
        "1. Philipp Lindenberger, Paul-Edouard Sarlin, Marc Pollefeys: “LightGlue: Local Feature Matching at Light Speed”, 2023; [arXiv:2306.13643](https://arxiv.org/abs/2306.13643).\n",
        "2. Zhengqi Li, Noah Snavely: “MegaDepth: Learning Single-View Depth Prediction from Internet Photos”, 2018; [arXiv:1804.00607](https://arxiv.org/abs/1804.00607).\n",
        "3. Jiaming Sun, Zehong Shen, Yuang Wang, Hujun Bao, Xiaowei Zhou: “LoFTR: Detector-Free Local Feature Matching with Transformers”, 2021; [arXiv:2104.00680](https://arxiv.org/abs/2104.00680)."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
